#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <ctime>

#include <cuda.h>

#include <GL/glut.h>

#include <sys/time.h>
#include <unistd.h>

#include <MLHI/ml_api.h>

#include "util.h"

struct vertex
{
    hrilFloat m;
    hrilFloat x[2];
    hrilFloat v[2];
    hrilFloat f[2];
    hrilFloat area_f[2];
    hrilFloat spring_f[2];
    hrilFloat damping_f[2];
};

struct triangle
{
    hrilFloat des_area;
    hrilFloat des_len[3];
    hrilFloat Vk;
    hrilFloat Sk;
    hrilFloat Sdk;
    int vertices[3];
};

// Constants
const hrilFloat gravity = 5;//9.8;
const hrilFloat outer_circle_radius = 1.5;
const hrilFloat outer_circle_stiffness = 1e4;

const int N = 8;
const int num_tris = 4*(N-1)*(N-1);
const int num_outer = 4*(N-1);
const hrilFloat initial_rotation = 0;//M_PI / 3;
const hrilFloat mass = 100.0;
const hrilFloat mass_per_point = mass / (N * N);

const hrilFloat Vk = 2.5e5;
const hrilFloat Sk = 5.0e3;//5e2;//1e3;
const hrilFloat Sdk = 3.8e1; //5.0;//10.0;

const hrilFloat haptics_circle_radius = 0.25;
const hrilFloat haptics_circle_stiffness = 100000.0;

// Interface variables
hrilFloat time_scale = 1.0;
hrilFloat time_step = 0.0010;//0.001;
bool playing = true;
bool draw_forces = true;
bool maglev_on = false;
ml_gain_vec_t gains;
ml_device_handle_t device_hdl;

// Host simulation variables
vertex vertices_h[N*N];
triangle triangles_h[num_tris];
int outer_vertices_h[num_outer];
hrilFloat haptics_circle_center_h[2] = { 0.0, 0.0 };
const hrilFloat outer_circle_center_h[2] = {0.2, 0.5};

// CUDA simulation variables
vertex* vertices_d = 0;
triangle* triangles_d = 0;
int* outer_vertices_d = 0;

int tick_callback(ml_device_handle_t dev_hdl, ml_position_t* pos)
{
    haptics_circle_center_h[0] = pos->values[0] * 250 ;
    haptics_circle_center_h[1] = pos->values[2] * 250 + .5;
    gains.values[0].ff = 0;
    gains.values[2].ff = 2;
    for (int i = 0; i < num_outer; i++)
    {
        hrilFloat penetration = haptics_circle_radius - dist(vertices_h[outer_vertices_h[i]].x, haptics_circle_center_h);

        if (penetration > 0)
        {
		gains.values[0].ff += .008 * (haptics_circle_center_h[0] - vertices_h[outer_vertices_h[i]].x[0]) * penetration * haptics_circle_stiffness / N;
		gains.values[2].ff += .008 * (haptics_circle_center_h[1] - vertices_h[outer_vertices_h[i]].x[1]) * penetration * haptics_circle_stiffness / N;
	}
    }
    ml_SetGainVecAxes(dev_hdl, ML_GAINSET_TYPE_NORMAL, gains);
    return ML_STATUS_OK;
}

bool init_mlhi()
{
	if (ml_Connect(&device_hdl, "192.168.0.2") != ML_STATUS_OK)
	{
		fprintf(stderr, "MLHI: Connection failed.\n");
		return false;
	}
	ml_fault_t fault;
	ml_GetFault(device_hdl, &fault);
	while (fault.value)
	{
		ml_ResetFault(device_hdl);
		usleep(4000);
		ml_GetFault(device_hdl, &fault);
	}
	ml_Takeoff(device_hdl);
	ml_GetGainVecAxes(device_hdl, ML_GAINSET_TYPE_NORMAL, &gains);
	gains.values[0].p = 10;
	gains.values[1].p = 2500;
	gains.values[2].p = 10;
	ml_SetGainVecAxes(device_hdl, ML_GAINSET_TYPE_NORMAL, gains);
	ml_RegisterCallbackTick(device_hdl, tick_callback);
	return true;
}

void cleanup_mlhi()
{
	ml_UnregisterCallbackTick(device_hdl);
	ml_Land(device_hdl);
	ml_Disconnect(device_hdl);
}

hrilFloat triangle_area_h(triangle* t)
{
	hrilFloat x1 = vertices_h[t->vertices[0]].x[0] - vertices_h[t->vertices[1]].x[0];
	hrilFloat y1 = vertices_h[t->vertices[0]].x[1] - vertices_h[t->vertices[1]].x[1];
	hrilFloat x2 = vertices_h[t->vertices[0]].x[0] - vertices_h[t->vertices[2]].x[0];
	hrilFloat y2 = vertices_h[t->vertices[0]].x[1] - vertices_h[t->vertices[2]].x[1];
	return (x1 * y2 - y1 * x2)/2.0;
}

__device__ hrilFloat triangle_area_d(vertex* vertices_d, triangle* t)
{
	hrilFloat x1 = vertices_d[t->vertices[0]].x[0] - vertices_d[t->vertices[1]].x[0];
	hrilFloat y1 = vertices_d[t->vertices[0]].x[1] - vertices_d[t->vertices[1]].x[1];
	hrilFloat x2 = vertices_d[t->vertices[0]].x[0] - vertices_d[t->vertices[2]].x[0];
	hrilFloat y2 = vertices_d[t->vertices[0]].x[1] - vertices_d[t->vertices[2]].x[1];
	return (x1 * y2 - y1 * x2)/2.0;
}

void init_scene()
{
	// init vertices
	for (int i = 0; i < N; i++)
	{
		for (int j = 0; j < N; j++)
		{
			vertices_h[N*i+j].m = mass_per_point;
			hrilFloat boxx = (hrilFloat)j / (N-1) - 0.5;
			hrilFloat boxy = (hrilFloat)i / (N-1) - 0.5;
			vertices_h[N*i+j].x[0] = cos(initial_rotation) * boxx - sin(initial_rotation) * boxy;
			vertices_h[N*i+j].x[1] = sin(initial_rotation) * boxx + cos(initial_rotation) * boxy + 0.5;
			vertices_h[N*i+j].v[0] = 0;
			vertices_h[N*i+j].v[1] = 0;
			vertices_h[N*i+j].f[0] = 0;
			vertices_h[N*i+j].f[1] = 0;
		}
	}

	// init triangles
	for (int i = 0; i < N-1; i++)
	{
		for (int j = 0; j < N-1; j++)
		{
			int hightriindex1 = 4 * ((N-1) * i + j);
			triangle* hightri1 = &triangles_h[hightriindex1];
			triangle* lowtri1 = &triangles_h[hightriindex1+1];
			triangle* hightri2 = &triangles_h[hightriindex1+2];
			triangle* lowtri2 = &triangles_h[hightriindex1+3];

			hightri1->Vk = Vk;
			hightri1->Sk = Sk;
			hightri1->Sdk = Sdk;

			lowtri1->Vk = Vk;
			lowtri1->Sk = Sk;
			lowtri1->Sdk = Sdk;

			hightri2->Vk = Vk;
			hightri2->Sk = Sk;
			hightri2->Sdk = Sdk;

			lowtri2->Vk = Vk;
			lowtri2->Sk = Sk;
			lowtri2->Sdk = Sdk;

			// CounterClockwise winding
			hightri1->vertices[0] = N*(i+1)+j;
			hightri1->vertices[1] = N*i+j+1;
			hightri1->vertices[2] = N*(i+1)+j+1;

			// CounterClockwise winding
			lowtri1->vertices[0] = N*i+j;
			lowtri1->vertices[1] = N*i+j+1;
			lowtri1->vertices[2] = N*(i+1)+j;

			// CounterClockwise winding
			hightri2->vertices[0] = N*(i+1)+j;
			hightri2->vertices[1] = N*i+j;
			hightri2->vertices[2] = N*(i+1)+j+1;

			// CounterClockwise winding
			lowtri2->vertices[0] = N*i+j;
			lowtri2->vertices[1] = N*i+j+1;
			lowtri2->vertices[2] = N*(i+1)+j+1;

			hightri1->des_area = triangle_area_h(hightri1);
			lowtri1->des_area  = triangle_area_h(lowtri1);
			hightri2->des_area = triangle_area_h(hightri2);
			lowtri2->des_area  = triangle_area_h(lowtri2);

			// desired displacements
			for (int k = 0; k < 3; k++)
			{
				hightri1->des_len[k] = dist(vertices_h[hightri1->vertices[k]].x, vertices_h[hightri1->vertices[(k+1)%3]].x);
				lowtri1->des_len[k] =  dist(vertices_h[lowtri1->vertices[k]].x,  vertices_h[lowtri1->vertices[(k+1)%3]].x);
				hightri2->des_len[k] = dist(vertices_h[hightri2->vertices[k]].x, vertices_h[hightri2->vertices[(k+1)%3]].x);
				lowtri2->des_len[k] =  dist(vertices_h[lowtri2->vertices[k]].x,  vertices_h[lowtri2->vertices[(k+1)%3]].x);
			}
		}
	}

	// init outer vertices
	for (int i = 0; i < N - 1; i++)
	{
		outer_vertices_h[4*i] = i;
		outer_vertices_h[4*i + 1] = N*(N-1)+N-1-i;
		outer_vertices_h[4*i + 2] = N*i+N-1;
		outer_vertices_h[4*i + 3] = N*(N-1-i);
	}

	// init device memory.
	if (!vertices_d)
		cudaMalloc(&vertices_d, sizeof(vertices_h));
	if (!triangles_d)
		cudaMalloc(&triangles_d, sizeof(triangles_h));
	if (!outer_vertices_d)
		cudaMalloc(&outer_vertices_d, sizeof(outer_vertices_h));

    cudaMemcpy(vertices_d, vertices_h, sizeof(vertices_h), cudaMemcpyHostToDevice);
    cudaMemcpy(triangles_d, triangles_h, sizeof(triangles_h), cudaMemcpyHostToDevice);
    cudaMemcpy(outer_vertices_d, outer_vertices_h, sizeof(outer_vertices_h), cudaMemcpyHostToDevice);
}

void calculate_forces()
{
    for (int i = 0; i < N*N; i++)
    {
        vertices_h[i].f[0] = 0;
        vertices_h[i].f[1] = -gravity * vertices_h[i].m; // apply gravity
        vertices_h[i].damping_f[0] = 0;
        vertices_h[i].damping_f[1] = 0;
        vertices_h[i].area_f[0] = 0;
        vertices_h[i].area_f[1] = 0;
        vertices_h[i].spring_f[0] = 0;
        vertices_h[i].spring_f[1] = 0;
    }

    // apply contact forces
    for (int i = 0; i < num_outer; i++)
    {
        hrilFloat penetration = dist(vertices_h[outer_vertices_h[i]].x, outer_circle_center_h) - outer_circle_radius;

        if (penetration > 0)
        {
            vertices_h[outer_vertices_h[i]].f[0] += (outer_circle_center_h[0] - vertices_h[outer_vertices_h[i]].x[0]) * penetration * outer_circle_stiffness;
            vertices_h[outer_vertices_h[i]].f[1] += (outer_circle_center_h[1] - vertices_h[outer_vertices_h[i]].x[1]) * penetration * outer_circle_stiffness;
        }

        if (maglev_on)
        {
            penetration = haptics_circle_radius - dist(vertices_h[outer_vertices_h[i]].x, haptics_circle_center_h);
    
            if (penetration > 0)
            {
                vertices_h[outer_vertices_h[i]].f[0] += (vertices_h[outer_vertices_h[i]].x[0] - haptics_circle_center_h[0]) * penetration * haptics_circle_stiffness;
                vertices_h[outer_vertices_h[i]].f[1] += (vertices_h[outer_vertices_h[i]].x[1] - haptics_circle_center_h[1]) * penetration * haptics_circle_stiffness;
            }
        }
    }

    // calculate internal forces
    for (int i = 0; i < num_tris; i++)
    {
        hrilFloat area = triangle_area_h(&triangles_h[i]);
        // Area compensation.
        if (area < 0.0)
        {
            printf("WARNING: Area < 0 for triangle %i\n", i);
        }

        hrilFloat area_err = area - triangles_h[i].des_area;
        hrilFloat center[2] = {
            (vertices_h[triangles_h[i].vertices[0]].x[0] 
           + vertices_h[triangles_h[i].vertices[1]].x[0] 
           + vertices_h[triangles_h[i].vertices[2]].x[0]) / 3.0,
            (vertices_h[triangles_h[i].vertices[0]].x[1] 
           + vertices_h[triangles_h[i].vertices[1]].x[1] 
           + vertices_h[triangles_h[i].vertices[2]].x[1]) / 3.0
        };

        for (int j = 0; j < 3; j++)
        {
            vertices_h[triangles_h[i].vertices[j]].area_f[0] += (center[0] - vertices_h[triangles_h[i].vertices[j]].x[0]) * area_err * triangles_h[i].Vk;
            vertices_h[triangles_h[i].vertices[j]].area_f[1] += (center[1] - vertices_h[triangles_h[i].vertices[j]].x[1]) * area_err * triangles_h[i].Vk;
        }

        // Spring compensation.
        for (int j = 0; j < 3; j++)
        {
            int prev = wrap(j-1,3);
            int next = wrap(j+1,3);

            hrilFloat len_prev = dist(vertices_h[triangles_h[i].vertices[prev]].x, vertices_h[triangles_h[i].vertices[j]].x);

            hrilFloat spr_err_prev = len_prev - triangles_h[i].des_len[prev];

            hrilFloat len_next = dist(vertices_h[triangles_h[i].vertices[next]].x, vertices_h[triangles_h[i].vertices[j]].x);

            hrilFloat spr_err_next =  len_next - triangles_h[i].des_len[j];

            vertices_h[triangles_h[i].vertices[j]].spring_f[0] += (vertices_h[triangles_h[i].vertices[prev]].x[0] - vertices_h[triangles_h[i].vertices[j]].x[0])/len_prev * spr_err_prev * triangles_h[i].Sk;
            vertices_h[triangles_h[i].vertices[j]].spring_f[1] += (vertices_h[triangles_h[i].vertices[prev]].x[1] - vertices_h[triangles_h[i].vertices[j]].x[1])/len_prev * spr_err_prev * triangles_h[i].Sk;

            hrilFloat v_proj_x = ((vertices_h[triangles_h[i].vertices[prev]].v[0] - vertices_h[triangles_h[i].vertices[j]].v[0])
                                * (vertices_h[triangles_h[i].vertices[prev]].x[0] - vertices_h[triangles_h[i].vertices[j]].x[0])
                                + (vertices_h[triangles_h[i].vertices[prev]].v[1] - vertices_h[triangles_h[i].vertices[j]].v[1])
                                * (vertices_h[triangles_h[i].vertices[prev]].x[1] - vertices_h[triangles_h[i].vertices[j]].x[1])) / len_prev;
            vertices_h[triangles_h[i].vertices[j]].damping_f[0] += v_proj_x * (vertices_h[triangles_h[i].vertices[prev]].x[0] - vertices_h[triangles_h[i].vertices[j]].x[0]) / len_prev * triangles_h[i].Sdk;
            vertices_h[triangles_h[i].vertices[j]].damping_f[1] += v_proj_x * (vertices_h[triangles_h[i].vertices[prev]].x[1] - vertices_h[triangles_h[i].vertices[j]].x[1]) / len_prev * triangles_h[i].Sdk;

            vertices_h[triangles_h[i].vertices[j]].spring_f[0] += (vertices_h[triangles_h[i].vertices[next]].x[0] - vertices_h[triangles_h[i].vertices[j]].x[0])/len_next * spr_err_next * triangles_h[i].Sk;
            vertices_h[triangles_h[i].vertices[j]].spring_f[1] += (vertices_h[triangles_h[i].vertices[next]].x[1] - vertices_h[triangles_h[i].vertices[j]].x[1])/len_next * spr_err_next * triangles_h[i].Sk;

            v_proj_x = ((vertices_h[triangles_h[i].vertices[next]].v[0] - vertices_h[triangles_h[i].vertices[j]].v[0])
                      * (vertices_h[triangles_h[i].vertices[next]].x[0] - vertices_h[triangles_h[i].vertices[j]].x[0])
                      + (vertices_h[triangles_h[i].vertices[next]].v[1] - vertices_h[triangles_h[i].vertices[j]].v[1])
                      * (vertices_h[triangles_h[i].vertices[next]].x[1] - vertices_h[triangles_h[i].vertices[j]].x[1])) / len_next;
            vertices_h[triangles_h[i].vertices[j]].damping_f[0] += v_proj_x * (vertices_h[triangles_h[i].vertices[next]].x[0] - vertices_h[triangles_h[i].vertices[j]].x[0]) / len_next *  triangles_h[i].Sdk;
            vertices_h[triangles_h[i].vertices[j]].damping_f[1] += v_proj_x * (vertices_h[triangles_h[i].vertices[next]].x[1] - vertices_h[triangles_h[i].vertices[j]].x[1]) / len_next *  triangles_h[i].Sdk;
        }
    }

    for (int i = 0; i < N*N; i++)
    {
        vertices_h[i].f[0] += vertices_h[i].area_f[0] + vertices_h[i].spring_f[0] + vertices_h[i].damping_f[0];
        vertices_h[i].f[1] += vertices_h[i].area_f[1] + vertices_h[i].spring_f[1] + vertices_h[i].damping_f[1];
    }
}

__global__ void init_and_contactf_d(vertex* vertices_d, int* outer_vertices_d, hrilFloat haptics_circle_center_d0, hrilFloat haptics_circle_center_d1, 
                                    hrilFloat outer_circle_center_d0, hrilFloat outer_circle_center_d1, bool maglev_on)
{
    int i = blockDim.x * blockIdx.x + threadIdx.x;
    
    if (i < N*N)
    {
        // Init forces
        vertices_d[i].f[0] = 0;
        vertices_d[i].f[1] = -gravity * vertices_d[i].m; // apply gravity
        vertices_d[i].damping_f[0] = 0;
        vertices_d[i].damping_f[1] = 0;
        vertices_d[i].area_f[0] = 0;
        vertices_d[i].area_f[1] = 0;
        vertices_d[i].spring_f[0] = 0;
        vertices_d[i].spring_f[1] = 0;

    }

    // Add contact forces
    if (i < num_outer)
    {
        hrilFloat penetration = sqrt((outer_circle_center_d0-vertices_d[outer_vertices_d[i]].x[0])*(outer_circle_center_d0-vertices_d[outer_vertices_d[i]].x[0])
                                   + (outer_circle_center_d1-vertices_d[outer_vertices_d[i]].x[1])*(outer_circle_center_d1-vertices_d[outer_vertices_d[i]].x[1]))
                                   - outer_circle_radius;

        if (penetration > 0)
        {
            vertices_d[outer_vertices_d[i]].f[0] += (outer_circle_center_d0 - vertices_d[outer_vertices_d[i]].x[0]) * penetration * outer_circle_stiffness;
            vertices_d[outer_vertices_d[i]].f[1] += (outer_circle_center_d1 - vertices_d[outer_vertices_d[i]].x[1]) * penetration * outer_circle_stiffness;
        }

        if (maglev_on)
        {
            penetration = haptics_circle_radius 
                              - sqrt((haptics_circle_center_d0-vertices_d[outer_vertices_d[i]].x[0])*(haptics_circle_center_d0-vertices_d[outer_vertices_d[i]].x[0])
                                   + (haptics_circle_center_d1-vertices_d[outer_vertices_d[i]].x[1])*(haptics_circle_center_d1-vertices_d[outer_vertices_d[i]].x[1]));

            if (penetration > 0)
            {
                vertices_d[outer_vertices_d[i]].f[0] += (vertices_d[outer_vertices_d[i]].x[0] - haptics_circle_center_d0) * penetration * haptics_circle_stiffness;
                vertices_d[outer_vertices_d[i]].f[1] += (vertices_d[outer_vertices_d[i]].x[1] - haptics_circle_center_d1) * penetration * haptics_circle_stiffness;
            }
        }
    }
}

__global__ void trianglef_d(vertex* vertices_d, triangle* triangles_d, int trii, int verti, int prev, int next)
{
    int i = 4*(blockDim.x * blockIdx.x + threadIdx.x) + trii;

    if (i < num_tris)
    {
        // Area compensation
        hrilFloat area_err = triangle_area_d(vertices_d, &triangles_d[i])- triangles_d[i].des_area;
        hrilFloat center[2] = {
            (vertices_d[triangles_d[i].vertices[0]].x[0] 
           + vertices_d[triangles_d[i].vertices[1]].x[0] 
           + vertices_d[triangles_d[i].vertices[2]].x[0]) / 3.0f,
            (vertices_d[triangles_d[i].vertices[0]].x[1] 
           + vertices_d[triangles_d[i].vertices[1]].x[1] 
           + vertices_d[triangles_d[i].vertices[2]].x[1]) / 3.0f
        };

        vertices_d[triangles_d[i].vertices[verti]].area_f[0] += (center[0] - vertices_d[triangles_d[i].vertices[verti]].x[0]) * area_err * triangles_d[i].Vk;
        vertices_d[triangles_d[i].vertices[verti]].area_f[1] += (center[1] - vertices_d[triangles_d[i].vertices[verti]].x[1]) * area_err * triangles_d[i].Vk;

        // Spring and damping compensation
        hrilFloat len_prev = dist(vertices_d[triangles_d[i].vertices[prev]].x, vertices_d[triangles_d[i].vertices[verti]].x);
        hrilFloat spr_err_prev = len_prev - triangles_d[i].des_len[prev];
        hrilFloat len_next = dist(vertices_d[triangles_d[i].vertices[next]].x, vertices_d[triangles_d[i].vertices[verti]].x);
        hrilFloat spr_err_next =  len_next - triangles_d[i].des_len[verti];

        // Trying to reduce recomputation, improve occupancy.
        hrilFloat diffXprev0 = vertices_d[triangles_d[i].vertices[prev]].x[0] - vertices_d[triangles_d[i].vertices[verti]].x[0];
        hrilFloat diffXprev1 = vertices_d[triangles_d[i].vertices[prev]].x[1] - vertices_d[triangles_d[i].vertices[verti]].x[1];
        hrilFloat diffXnext0 = vertices_d[triangles_d[i].vertices[next]].x[0] - vertices_d[triangles_d[i].vertices[verti]].x[0];
        hrilFloat diffXnext1 = vertices_d[triangles_d[i].vertices[next]].x[1] - vertices_d[triangles_d[i].vertices[verti]].x[1];

        hrilFloat err_sdk_over_prev = spr_err_prev * triangles_d[i].Sk / len_prev;
        hrilFloat err_sdk_over_next = spr_err_next * triangles_d[i].Sk / len_next;

        vertices_d[triangles_d[i].vertices[verti]].spring_f[0] += diffXprev0 * err_sdk_over_prev;
        vertices_d[triangles_d[i].vertices[verti]].spring_f[1] += diffXprev1 * err_sdk_over_prev;

        hrilFloat v_proj_x = ((vertices_d[triangles_d[i].vertices[prev]].v[0] - vertices_d[triangles_d[i].vertices[verti]].v[0]) * diffXprev0
                            + (vertices_d[triangles_d[i].vertices[prev]].v[1] - vertices_d[triangles_d[i].vertices[verti]].v[1]) * diffXprev1) / len_prev;
        vertices_d[triangles_d[i].vertices[verti]].damping_f[0] += v_proj_x * diffXprev0 / len_prev * triangles_d[i].Sdk;
        vertices_d[triangles_d[i].vertices[verti]].damping_f[1] += v_proj_x * diffXprev1 / len_prev * triangles_d[i].Sdk;

        vertices_d[triangles_d[i].vertices[verti]].spring_f[0] += diffXnext0 * err_sdk_over_next;
        vertices_d[triangles_d[i].vertices[verti]].spring_f[1] += diffXnext1 * err_sdk_over_next;

        v_proj_x = ((vertices_d[triangles_d[i].vertices[next]].v[0] - vertices_d[triangles_d[i].vertices[verti]].v[0]) * diffXnext0
                  + (vertices_d[triangles_d[i].vertices[next]].v[1] - vertices_d[triangles_d[i].vertices[verti]].v[1]) * diffXnext1) / len_next;
        vertices_d[triangles_d[i].vertices[verti]].damping_f[0] += v_proj_x * diffXnext0 / len_next *  triangles_d[i].Sdk;
        vertices_d[triangles_d[i].vertices[verti]].damping_f[1] += v_proj_x * diffXnext1 / len_next *  triangles_d[i].Sdk;

    }
}

__global__ void integrate_d(vertex* vertices_d, float dt)
{
    int i = blockDim.x * blockIdx.x + threadIdx.x;

    if (i < N*N)
    {
        // Add up diff forces
        vertices_d[i].f[0] += vertices_d[i].area_f[0] + vertices_d[i].spring_f[0] + vertices_d[i].damping_f[0];
        vertices_d[i].f[1] += vertices_d[i].area_f[1] + vertices_d[i].spring_f[1] + vertices_d[i].damping_f[1];

        // Perform integration
        vertices_d[i].x[0] += vertices_d[i].v[0] * dt;
        vertices_d[i].x[1] += vertices_d[i].v[1] * dt;
        vertices_d[i].v[0] += vertices_d[i].f[0] / vertices_d[i].m * dt;
        vertices_d[i].v[1] += vertices_d[i].f[1] / vertices_d[i].m * dt;
    }
}

void integrate(float dt)
{
    const int blocksize = 16;//N*N/2 + N*N%2;
    const int blocks = 96;//N*N/blocksize + !!(N*N%blocksize);
    init_and_contactf_d<<<blocks, blocksize>>>(vertices_d, outer_vertices_d, haptics_circle_center_h[0], haptics_circle_center_h[1], 
                                               outer_circle_center_h[0], outer_circle_center_h[1], maglev_on);
    for (int i = 0; i < 4; i++)
    {
        trianglef_d<<<blocks, blocksize>>>(vertices_d, triangles_d, i, 0, 2, 1);
        trianglef_d<<<blocks, blocksize>>>(vertices_d, triangles_d, i, 1, 0, 2);
        trianglef_d<<<blocks, blocksize>>>(vertices_d, triangles_d, i, 2, 1, 0);
    }
    integrate_d<<<blocks, blocksize>>>(vertices_d, dt);

    cudaMemcpy(vertices_h, vertices_d, sizeof(vertices_h), cudaMemcpyDeviceToHost);
}

void draw_scene()
{
    glClear(GL_COLOR_BUFFER_BIT);

    glBegin(GL_TRIANGLES);
    for (int i = 0; i < num_tris; i++)
    {
        glColor4f(0.0, (hrilFloat)i / num_tris, 1.0, 1.0);
        glVertex3f(vertices_h[triangles_h[i].vertices[0]].x[0],
                vertices_h[triangles_h[i].vertices[0]].x[1],
                0.0);
        glVertex3f(vertices_h[triangles_h[i].vertices[1]].x[0],
                vertices_h[triangles_h[i].vertices[1]].x[1],
                0.0);
        glVertex3f(vertices_h[triangles_h[i].vertices[2]].x[0],
                vertices_h[triangles_h[i].vertices[2]].x[1],
                0.0);
    }
    glEnd();

    if (draw_forces)
    {
        glBegin(GL_LINES);
        for (int i = 0; i < N; i++)
        {
            for (int j = 0; j < N; j++)
            {
                glColor3f(1.0, 0.0, 0.0);
                glVertex3f(vertices_h[N*i+j].x[0],
                        vertices_h[N*i+j].x[1],
                        0.0);
                glVertex3f(vertices_h[N*i+j].x[0] + 0.01 * vertices_h[N*i+j].f[0],
                        vertices_h[N*i+j].x[1] + 0.01 * vertices_h[N*i+j].f[1],
                        0.0);
                
                glColor3f(1.0, 0.0, 1.0);
                glVertex3f(vertices_h[N*i+j].x[0],
                        vertices_h[N*i+j].x[1],
                        0.0);
                glVertex3f(vertices_h[N*i+j].x[0] + 0.01 * vertices_h[N*i+j].area_f[0],
                        vertices_h[N*i+j].x[1] + 0.01 * vertices_h[N*i+j].area_f[1],
                        0.0);
                glColor3f(1.0, 1.0, 0.0);
                glVertex3f(vertices_h[N*i+j].x[0],
                        vertices_h[N*i+j].x[1],
                        0.0);
                glVertex3f(vertices_h[N*i+j].x[0] + 0.01 * vertices_h[N*i+j].damping_f[0],
                        vertices_h[N*i+j].x[1] + 0.01 * vertices_h[N*i+j].damping_f[1],
                        0.0);
                glColor3f(1.0, 1.0, 1.0);
                glVertex3f(vertices_h[N*i+j].x[0],
                        vertices_h[N*i+j].x[1],
                    0.0);
                glVertex3f(vertices_h[N*i+j].x[0] + 0.01 * vertices_h[N*i+j].spring_f[0],
                        vertices_h[N*i+j].x[1] + 0.01 * vertices_h[N*i+j].spring_f[1],
                        0.0);
                /*
                glColor3f(0.0, 0.0, 0.0);
                glVertex3f(vertices_h[N*i+j].x[0],
                        vertices_h[N*i+j].x[1],
                    0.0);
                glVertex3f(vertices_h[N*i+j].x[0] + 1.0 * vertices_h[N*i+j].v[0],
                        vertices_h[N*i+j].x[1] + 1.0 * vertices_h[N*i+j].v[1],
                        0.0);
                */
            }
        }
        glEnd();
    }

    glColor3f(0.0, 1.0, 0.0);
    glBegin(GL_LINE_LOOP);
    for (int i = 0; i < 2*36; i++)
    {
        glVertex3f(outer_circle_center_h[0] + outer_circle_radius * sin(i * 5 * 3.14159 / 180),
                outer_circle_center_h[1] + outer_circle_radius * cos(i * 5 * 3.14159 / 180),
                0.0);
    }
    glEnd();

    if (maglev_on)
    {
        glColor3f(0.0, 1.0, 0.0);
        glBegin(GL_LINE_LOOP);
        for (int i = 0; i < 2*36; i++)
        {
            glVertex3f(haptics_circle_center_h[0] + haptics_circle_radius * sin(i * 5 * 3.14159 / 180),
                    haptics_circle_center_h[1] + haptics_circle_radius * cos(i * 5 * 3.14159 / 180),
                    0.0);
        }
        glEnd();
    }
 
    glutSwapBuffers();
}

void display()
{
    static double last_draw_time = get_time();
    static double last_integrate_time = get_time();

    static int frames = 0;
    static int integrations = 0;
    static int prev_time = time(NULL);
    int cur_time = time(NULL);
    if (cur_time != prev_time)
    {
        printf("FPS: %i, IPS: %i\n", frames, integrations);
        prev_time = cur_time;
        frames = 0;
        integrations = 0;
    }

    int integration_counts = 0;
    while (get_time() - last_integrate_time > time_scale * (1.0/60.0) * time_step)
    {
        if (playing)
        {
            integrate(time_step);
            integrations++;
            last_integrate_time += time_step/time_scale;
            if (++integration_counts > 3 * time_scale * (1.0/60.0)/(time_step))
            {
                //printf("WARNING: not running real time.\n");
                last_integrate_time = get_time();
                break;
            }
        }
        else
        {
            last_integrate_time = get_time();
        }
    }
    

    //if (get_time() - last_draw_time > (1.0/60.0))
    {
        //usleep(100000*(get_time()-last_draw_time));
        //last_draw_time = get_time();
        draw_scene();
        frames++;
    }

    glutPostRedisplay();
}

void keyboard(unsigned char key, int x, int y)
{
    if (key == 'q')
        exit(0);

    if (key == '+')
    {
        time_scale += 0.1;
        printf("time scale: %f\n", time_scale);
    }
    else if (key == '-')
    {
        if (time_scale - 0.1 > 0.01)
            time_scale -= 0.1;
        printf("time scale: %f\n", time_scale);
    }

    else if (key == '(')
    {
        if (time_step - 0.0001 > 0.00001)
            time_step -= 0.0001;
        printf("time step: %f\n", time_step);
    }
    else if (key == ')')
    {
        time_step += 0.0001;
        printf("time step: %f\n", time_step);
    }

    else if (key == 'r')
    {
        init_scene();
        printf("Restarted.\n");
    }

    else if (key == ' ')
    {
        playing = !playing;
        printf("Playing %s.\n", playing ? "on" : "off");
    }

    else if (key == 'n' && !playing)
    {
        integrate(time_step);
        printf("Step.\n");
    }

    else if (key == 'f')
    {
        draw_forces = !draw_forces;
        printf("Drawing forces %s.\n", draw_forces ? "on" : "off");
    }

    else if (key == 'm')
    {
        if (!maglev_on)
        {
            printf("Initializing maglev.\n");
            maglev_on = init_mlhi();
        }
        else
        {
            maglev_on = false;
            printf("Shutting down maglev.\n");
            cleanup_mlhi();
        }
    }
}

void cleanup()
{
    if (maglev_on)
        cleanup_mlhi();
    cudaFree((void**)vertices_d);
    cudaFree((void**)triangles_d);
    cudaFree((void**)outer_vertices_d);
    printf("Quitting.\n");
}

void init_glut(int* argc, char** argv)
{
    glutInit(argc, argv);
    glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB);
    glutInitWindowSize(1000, 1000);
    glutInitWindowPosition(100, 100);
    glutCreateWindow("2D Volume Conserved Soft Body");
    glutDisplayFunc(display);
    glutKeyboardFunc(keyboard);
    atexit(cleanup);

}

int main(int argc, char** argv)
{
    init_scene();
    printf("Scene inited.\n");
    init_glut(&argc, argv);
    printf("GLUT inited.\n");

    glutMainLoop();
}
